#pragma once

#include <iostream>
#include <streambuf>

// 1. 自定义的 streambuf，用于在每行开头添加制表符
class TabbedStreambuf : public std::streambuf {
public:
    // 构造函数接收原始的 streambuf 指针
    TabbedStreambuf(std::streambuf* original_buf)
        : original_buf(original_buf), is_at_start_of_line(true) {}

protected:
    // 2. 重写 overflow()，这是处理输出的核心
    // 当缓冲区满或遇到特定字符（如 std::endl）时被调用
    virtual int_type overflow(int_type c = traits_type::eof()) override {
        if (traits_type::eq_int_type(c, traits_type::eof())) {
            // 如果是文件结束符，直接同步原始缓冲区
            return original_buf->pubsync() == -1 ? traits_type::eof() : traits_type::not_eof(c);
        }

        // 3. 检查是否在行首
        if (is_at_start_of_line) {
            // 如果在行首，先输出一个制表符
            if (traits_type::eq_int_type(original_buf->sputc('\t'), traits_type::eof())) {
                return traits_type::eof(); // 如果输出失败，返回 EOF
            }
            is_at_start_of_line = false; // 输出后，就不再是行首了
        }

        // 输出实际的字符
        if (traits_type::eq_int_type(original_buf->sputc((char_type)c), traits_type::eof())) {
            return traits_type::eof(); // 如果输出失败，返回 EOF
        }

        // 4. 检查输出的字符是否是换行符
        if (traits_type::eq_int_type(c, traits_type::to_int_type('\n'))) {
            is_at_start_of_line = true; // 如果是换行符，下一行就是行首
        }

        return traits_type::not_eof(c);
    }
    
    // 也需要重写 sync() 以便刷新缓冲区
    virtual int sync() override {
        return original_buf->pubsync();
    }

private:
    std::streambuf* original_buf; // 指向原始的 cout streambuf
    bool is_at_start_of_line;     // 状态标志
};

// 5. RAII 管理器，用于自动替换和恢复 streambuf
class CoutTabber {
public:
    CoutTabber() : original_cout_buf(std::cout.rdbuf()), tabbed_buf(original_cout_buf) {
        // 构造时，用我们的 TabbedStreambuf 替换 std::cout 的 streambuf
        std::cout.rdbuf(&tabbed_buf);
    }

    ~CoutTabber() {
        // 析构时，恢复原始的 streambuf
        std::cout.rdbuf(original_cout_buf);
    }

private:
    std::streambuf* original_cout_buf;
    TabbedStreambuf tabbed_buf;
};
